package me.zk.aspect;

import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import me.zk.annotation.Limit;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.Map;

/**
 * @author zhangkai
 * @date 2023/6/30 17:42
 */
@Slf4j
@Aspect
@Component
public class LimitAspect {

    private final Map<String, RateLimiter> limiterMap = Maps.newConcurrentMap();

    @Around("@annotation(me.zk.annotation.Limit)")
    public Object around(ProceedingJoinPoint pjp) throws Throwable {
        MethodSignature signature = (MethodSignature) pjp.getSignature();
        Method method = signature.getMethod();

        Limit limit = method.getAnnotation(Limit.class);
        if (limit != null) {
            String key = limit.key();
            RateLimiter rateLimiter;
            if (!limiterMap.containsKey(key)) {
                rateLimiter = RateLimiter.create(limit.permitsPerSecond());
                limiterMap.put(key,rateLimiter);
                log.info("新建了令牌桶={},容量={}",key,limit.permitsPerSecond());
            }
            rateLimiter = limiterMap.get(key);
            boolean acquire = rateLimiter.tryAcquire(limit.timeout(),limit.timeunit());
            if (!acquire) {
                log.info("令牌桶={},未获取",key);
                throw new RuntimeException(limit.msg());
            }
        }


        return pjp.proceed();

    }

}
